package io.gitee.zhangbinhub.acp.cloud.oauth.authentication;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.gitee.zhangbinhub.acp.cloud.oauth.domain.SecurityUserDetailsService;
import io.gitee.zhangbinhub.acp.cloud.oauth.vo.TokenUserInfoVo;
import io.gitee.zhangbinhub.acp.cloud.resource.server.constant.AcpCloudResourceServerConstant;
import org.bouncycastle.util.encoders.Base64;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;

import java.security.Principal;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Collectors;

public class OAuth2UserPasswordAuthenticationProvider implements AuthenticationProvider {
    private static final String ERROR_URI = "https://datatracker.ietf.org/doc/html/rfc6749#section-5.2";
    private final ObjectMapper objectMapper;
    private final SecurityUserDetailsService securityUserDetailsService;
    private final OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator;
    private final OAuth2AuthorizationService oAuth2AuthorizationService;

    public OAuth2UserPasswordAuthenticationProvider(ObjectMapper objectMapper, SecurityUserDetailsService securityUserDetailsService, OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator, OAuth2AuthorizationService oAuth2AuthorizationService) {
        this.objectMapper = objectMapper;
        this.securityUserDetailsService = securityUserDetailsService;
        this.tokenGenerator = tokenGenerator;
        this.oAuth2AuthorizationService = oAuth2AuthorizationService;
    }

    @Override
    public Authentication authenticate(Authentication authentication) throws AuthenticationException {
        OAuth2UserPasswordAuthenticationToken userPasswordAuthenticationToken = (OAuth2UserPasswordAuthenticationToken) authentication;
        String username = (String) userPasswordAuthenticationToken.getPrincipal();
//        String password = (String) userPasswordAuthenticationToken.getCredentials();
        UserDetails userDetails = securityUserDetailsService.loadUserByUsername(username);
        // 验证用户身份，自行添加 start
        // 验证用户身份，自行添加 end
        OAuth2ClientAuthenticationToken clientPrincipal = (OAuth2ClientAuthenticationToken) userPasswordAuthenticationToken.getClientPrincipal();
        RegisteredClient registeredClient = clientPrincipal.getRegisteredClient();
        if (registeredClient == null) {
            OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.UNAUTHORIZED_CLIENT,
                    "The client unauthorized.", ERROR_URI);
            throw new OAuth2AuthenticationException(error);
        }
        OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.withRegisteredClient(registeredClient)
                .authorizationGrantType(new AuthorizationGrantType("password"))
                .attributes(attrs -> attrs.put(Principal.class.getName(), clientPrincipal))
                .authorizedScopes(registeredClient.getScopes())
                .principalName(username);
        DefaultOAuth2TokenContext.Builder tokenContextBuilder = DefaultOAuth2TokenContext.builder()
                .registeredClient(registeredClient)
                .principal(userPasswordAuthenticationToken)
                .authorizationServerContext(AuthorizationServerContextHolder.getContext())
                .authorizedScopes(registeredClient.getScopes())
                .authorizationGrantType(new AuthorizationGrantType("password"))
                .authorizationGrant(userPasswordAuthenticationToken);

        boolean authenticated = false;
        // ----- Access token -----
        OAuth2TokenContext tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.ACCESS_TOKEN).build();
        OAuth2Token generatedAccessToken = this.tokenGenerator.generate(tokenContext);
        if (generatedAccessToken == null) {
            OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
                    "The token generator failed to generate the access token.", ERROR_URI);
            throw new OAuth2AuthenticationException(error);
        }
        OAuth2AccessToken accessToken = new OAuth2AccessToken(OAuth2AccessToken.TokenType.BEARER,
                generatedAccessToken.getTokenValue(), generatedAccessToken.getIssuedAt(),
                generatedAccessToken.getExpiresAt(), tokenContext.getAuthorizedScopes());
        if (generatedAccessToken instanceof ClaimAccessor) {
            Map<String, Object> claims = new HashMap<>(((ClaimAccessor) generatedAccessToken).getClaims());
            claims.put(AcpCloudResourceServerConstant.TOKEN_CLAIMS_AUTHORITIES, userDetails.getAuthorities().stream().map(GrantedAuthority::getAuthority).collect(Collectors.toSet()));
            // 添加附加的用户信息 start
            TokenUserInfoVo userInfo = new TokenUserInfoVo();
            userInfo.setAppId("app id");
            userInfo.setId("id");
            userInfo.setLoginNo("login no");
            userInfo.setName("user name");
            userInfo.setMobile("13000000000");
            userInfo.setLoginTime(1603237865123L);
            try {
                claims.put(AcpCloudResourceServerConstant.TOKEN_CLAIMS_USER_INFO, Base64.toBase64String(objectMapper.writeValueAsBytes(userInfo)));
            } catch (JsonProcessingException e) {
                throw new RuntimeException(e);
            }
            // 添加附加的用户信息 end
            authorizationBuilder.token(accessToken, (metadata) ->
                    metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, claims));
            if (!userDetails.getAuthorities().isEmpty() || !accessToken.getScopes().isEmpty()) {
                authenticated = true;
            }
        } else {
            authorizationBuilder.accessToken(accessToken);
        }

        // ----- Refresh token -----
        OAuth2RefreshToken refreshToken = null;
        if (registeredClient.getAuthorizationGrantTypes().contains(AuthorizationGrantType.REFRESH_TOKEN) &&
                !clientPrincipal.getClientAuthenticationMethod().equals(ClientAuthenticationMethod.NONE)) {
            tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.REFRESH_TOKEN).build();
            OAuth2Token generatedRefreshToken = this.tokenGenerator.generate(tokenContext);
            if (!(generatedRefreshToken instanceof OAuth2RefreshToken)) {
                OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
                        "The token generator failed to generate the refresh token.", ERROR_URI);
                throw new OAuth2AuthenticationException(error);
            }
            refreshToken = (OAuth2RefreshToken) generatedRefreshToken;
            authorizationBuilder.refreshToken(refreshToken);
        }
        OAuth2Authorization oAuth2Authorization = authorizationBuilder.build();
        oAuth2AuthorizationService.save(oAuth2Authorization);
        OAuth2AccessTokenAuthenticationToken authenticationToken = new OAuth2AccessTokenAuthenticationToken(
                registeredClient,
                clientPrincipal,
                accessToken,
                refreshToken,
                userPasswordAuthenticationToken.getAdditionalParameters());
        authenticationToken.setAuthenticated(authenticated);
        return authenticationToken;
    }

    @Override
    public boolean supports(Class<?> authentication) {
        return OAuth2UserPasswordAuthenticationToken.class.isAssignableFrom(authentication);
    }
}
